#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t ALIGN_NUM = 32;  
constexpr int32_t BUFFER_NUM = 1;   
class KernelDepthToSpace_1{
public:
    __aicore__ inline KernelDepthToSpace_1() {} 
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, 
                                uint32_t blockLength,
                                uint32_t N, uint32_t H, uint32_t W, uint32_t C, uint32_t B,
                                uint32_t baseW)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->N = N;
        this->H = H;
        this->W = W;   
        this->C = C;
        this->B = B;
        this->baseW = baseW;
        xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x, blockLength);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, blockLength);
        pipe.InitBuffer(Que, BUFFER_NUM, (baseW * (this->C) * sizeof(DTYPE_X)));
    } 
    __aicore__ inline void Process()
    { 
        uint32_t block_id = GetBlockIdx();
        uint32_t core_num = GetBlockNum();
        uint32_t calcWcnt;
        uint32_t NH = (this->N) * (this->H);
        uint16_t blockCnt;
        uint16_t blockLen;
        uint16_t srcStride;   
        uint32_t C_B = C/B;
        for(int outloop = 0; outloop < NH; outloop++) {
            uint32_t x_out_offset = outloop * (this->W) * (this->C);
            calcWcnt = baseW;
            for(int inloop = block_id * baseW; inloop <(this->W); inloop += core_num * baseW)
            {
                calcWcnt = baseW;
                uint32_t x_offset = x_out_offset + inloop * (this->C);
                if((inloop + baseW) > (this->W)) {
                    calcWcnt = (this->W) - inloop;
                }
                LocalTensor<DTYPE_X> dLocal = Que.AllocTensor<DTYPE_X>();  
                uint32_t calcLength = calcWcnt * (this->C);
                DataCopy(dLocal, xGm[x_offset], calcLength);
                Que.EnQue(dLocal);
                LocalTensor<DTYPE_X> d1Local = Que.DeQue<DTYPE_X>();
                for(int j=0; j<(this->B); j++)
                {
                    uint32_t d_offset = j * (C_B) ;  
                    uint32_t y_offset = x_out_offset + (d_offset * W) + inloop * (C_B);
                   blockCnt = calcWcnt;
                   blockLen = (this->C) / (this->B) * sizeof(DTYPE_X)  / 32 ;
                   srcStride = blockLen * (this->B - 1);                    
                   if(blockCnt == 1)
                   {
                       srcStride = 0;
                   }
                   DataCopyParams copyParams{blockCnt, blockLen, srcStride, 0};
                   DataCopy(yGm[y_offset], d1Local[d_offset], copyParams);
                }
                Que.FreeTensor(d1Local);
            }
        }
    }
private:
    TPipe pipe;
    TQueBind<TPosition::VECIN, TPosition::VECOUT, BUFFER_NUM> Que;
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_Y> yGm;
    uint32_t N, H, W, C, B, baseW;
};
class KernelDepthToSpace_5{
public:
    __aicore__ inline KernelDepthToSpace_5() {} 
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, 
                                uint32_t blockLength,
                                uint32_t N, uint32_t H, uint32_t W, uint32_t C, uint32_t B,
                                uint32_t baseW)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->N = N;
        this->H = H;
        this->W = W;   
        this->C = C;
        this->B = B;
        this->baseW = baseW;
        xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x, blockLength);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, blockLength);
        this->elemPerBlock = (ALIGN_NUM / sizeof(DTYPE_X));
        this->C_B = C / B;
        this->C_BAlign = RoundUp(C_B, elemPerBlock);            
        pipe.InitBuffer(Que, BUFFER_NUM, (baseW * (B * C_BAlign) * sizeof(DTYPE_X)));
    } 
    __aicore__ inline void Process()
    { 
        uint32_t block_id = GetBlockIdx();
        uint32_t core_num = GetBlockNum();
        uint32_t calcWcnt;
        uint32_t NH = (this->N) * (this->H);
        uint16_t blockCnt;
        uint16_t blockLen;
        uint16_t srcStride;      
        for(int outloop = 0; outloop < NH; outloop++) {
            uint32_t x_out_offset = outloop * (this->W) * (this->C);
            calcWcnt = baseW;
            for(int inloop = block_id * baseW; inloop <(this->W); inloop += core_num * baseW)
            {
                calcWcnt = baseW;
                uint32_t x_offset = x_out_offset + inloop * (this->C);
                if((inloop + baseW) > (this->W)) {
                    calcWcnt = (this->W) - inloop;
                }
                LocalTensor<DTYPE_X> dLocal = Que.AllocTensor<DTYPE_X>();  
                uint16_t x_blockCount  = calcWcnt * (this->B);
                uint32_t x_blockLen = (this->C) / (this->B) * sizeof(DTYPE_X);
                DataCopyExtParams copyParams{x_blockCount, x_blockLen, 0, 0, 0};
                DataCopyPadExtParams<DTYPE_X> padParams{true, 0, 0, 0}; 
                DataCopyPad(dLocal, xGm[x_offset], copyParams, padParams);                 
                Que.EnQue(dLocal);
                LocalTensor<DTYPE_X> d1Local = Que.DeQue<DTYPE_X>();
                for(int j=0; j<(this->B); j++)
                {
                    uint32_t y_offset = x_out_offset + j * (C_B * W) + inloop * (C_B);
                    uint32_t d_offset = j * (this->C_BAlign) ;                    
                   uint16_t y_blockCnt = calcWcnt;
                   uint32_t y_blockLen = this->C /this->B * sizeof(DTYPE_Y);
                   uint32_t y_srcStride = y_blockLen * (this->B - 1);
                   if(y_blockCnt == 1)
                   {
                       y_srcStride = 0;
                   }
                    DataCopyExtParams copyParams{y_blockCnt, y_blockLen, y_srcStride, 0, 0};
                    DataCopyPad(yGm[y_offset], dLocal[d_offset], copyParams);                     
                }
                Que.FreeTensor(d1Local);
            }
        }
    }
    __aicore__ inline uint32_t RoundUp(uint32_t a, uint32_t b)
    { 
        if (b == 0) {
            return 0;
        }
        return ((a + b - 1) / b) * b;
    }  
private:
    TPipe pipe;
    TQueBind<TPosition::VECIN, TPosition::VECOUT, BUFFER_NUM> Que;
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_Y> yGm;
    uint32_t N, H, W, C, B, baseW;
    uint32_t C_BAlign;
    uint32_t C_B;
    uint32_t elemPerBlock;
};
class KernelDepthToSpace_3{
public:
    __aicore__ inline KernelDepthToSpace_3() {} 
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, 
                                uint32_t blockLength,
                                uint32_t N, uint32_t H, uint32_t W, uint32_t C, uint32_t B,
                                uint32_t baseW)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->N = N;
        this->H = H;
        this->W = W;   
        this->C = C;
        this->B = B;
        xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x, blockLength);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, blockLength);
        uint32_t elemPerBlock = (ALIGN_NUM / sizeof(DTYPE_X));
        this->WC = (this->W) * (this->C);
        this->WCAlign = RoundUp(WC, elemPerBlock);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, (WCAlign * sizeof(DTYPE_X)));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, (WCAlign * sizeof(DTYPE_Y)));
    } 
    __aicore__ inline void Process()
    { 
        uint32_t block_id = GetBlockIdx();
        uint32_t core_num = GetBlockNum();
        this->x_blockCount = 1;
        this->x_blockLen  = this->WC * sizeof(DTYPE_X);
        this->x_dstStride = 0;
        this->y_blockLen = this->WC * sizeof(DTYPE_Y);
        for(int outloop = 0; outloop < (this->N); outloop++) {
            uint32_t x_out_offset = outloop * (this->C) * (this->H) * (this->W);
            for(int inloop = block_id; inloop < (this->H) ; inloop += core_num)
            {   
                this->x_offset = x_out_offset + inloop * (this->W) * (this->C); 
                this->y_offset = this->x_offset;
                CopyIn();
                Compute();
                CopyOut();
            }
        }
    }
private:
    __aicore__ inline uint32_t RoundUp(uint32_t a, uint32_t b)
    { 
        if (b == 0) {
            return 0;
        }
        return ((a + b - 1) / b) * b;
    }  
    __aicore__ inline void CopyIn()
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        DataCopyExtParams copyParams{x_blockCount, x_blockLen, x_dstStride, 0, 0};
        DataCopyPadExtParams<DTYPE_X> padParams{true, 0, 0, 0}; 
        DataCopyPad(xLocal, xGm[this->x_offset], copyParams, padParams); 
        inQueueX.EnQue(xLocal);
    }    
    __aicore__ inline void Compute() 
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();
        for(int i=0; i<(this->W); i++)
        {
            for(int j=0; j<(this->C); j++)
            {
                uint64_t y_index =  ((j % (B * B))/B) * ((C / B ) * W ) + (i * B + j % B) * (C/(B * B)) + j/(B * B);
                uint64_t x_index = i * (this->C) + j;
                yLocal.SetValue((y_index), xLocal.GetValue(x_index));
            }
        }        
        outQueueY.EnQue(yLocal);
        inQueueX.FreeTensor(xLocal);
    }
    __aicore__ inline void CopyOut()
    {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        DataCopyExtParams copyParams{1, y_blockLen, 0, 0, 0};
        DataCopyPad(yGm[this->y_offset], yLocal, copyParams); 
        outQueueY.FreeTensor(yLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_Y> yGm;
    uint32_t N, H, W, C, B;
    uint32_t WCAlign;
    uint32_t WC;
    uint16_t x_blockCount;
    uint32_t x_blockLen;
    uint32_t x_srcStride;
    uint32_t x_dstStride;
    uint32_t y_blockLen;
    uint32_t x_offset;
    uint32_t y_offset;
};
class KernelDepthToSpace_2{
public:
    __aicore__ inline KernelDepthToSpace_2() {} 
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, 
                                uint32_t blockLength,
                                uint32_t N, uint32_t H, uint32_t W, uint32_t C, uint32_t B,
                                uint32_t baseW)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->N = N;
        this->H = H;
        this->W = W;   
        this->C = C;
        this->B = B;
        xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x, blockLength);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, blockLength);
        uint32_t elemPerBlock = (ALIGN_NUM / sizeof(DTYPE_X));
        this->WAlign = RoundUp(W, elemPerBlock);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, (WAlign * (this->B) * sizeof(DTYPE_X)));
        uint32_t WBAlign = RoundUp(W*B, elemPerBlock);
        pipe.InitBuffer(outQueueY, BUFFER_NUM, (WBAlign * sizeof(DTYPE_X)));
    } 
    __aicore__ inline void Process()
    { 
        uint32_t block_id = GetBlockIdx();
        uint32_t core_num = GetBlockNum();
        this->x_blockCount = B;
        this->x_blockLen  = W * sizeof(DTYPE_X);
        this->x_dstStride = (H * (C / B / B) - 1) * W * sizeof(DTYPE_X);
        this->y_blockLen = B * W * sizeof(DTYPE_Y);
        for(int outloop = 0; outloop < (this->N); outloop++) {
            uint32_t x_out_offset = outloop * (this->C) * (this->H) * (this->W);
            for(int inloop = block_id; inloop < (((this->C) / (this->B) * (this->H))) ; inloop += core_num)
            {   
                this->x_offset = x_out_offset + ((inloop * H * (C / B))% (C * H) + (inloop /  (B))) * W; 
                this->y_offset = x_out_offset + inloop * (B * W);
                CopyIn();
                Compute();
                CopyOut();
            }
        }
    }
private:
    __aicore__ inline uint32_t RoundUp(uint32_t a, uint32_t b)
    { 
        if (b == 0) {
            return 0;
        }
        return ((a + b - 1) / b) * b;
    }  
    __aicore__ inline void CopyIn()
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        DataCopyExtParams copyParams{x_blockCount, x_blockLen, x_dstStride, 0, 0};
        DataCopyPadExtParams<DTYPE_X> padParams{true, 0, 0, 0}; 
        DataCopyPad(xLocal, xGm[this->x_offset], copyParams, padParams); 
        inQueueX.EnQue(xLocal);
    }    
    __aicore__ inline void Compute() 
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();
        for(int i = 0; i<(this->W); i++)
        {
            for(int j=0; j<(this->B); j++)
            {
                yLocal.SetValue((i * (this->B) + j), xLocal.GetValue(j*(this->WAlign) + i));
            }
        }
        outQueueY.EnQue(yLocal);
        inQueueX.FreeTensor(xLocal);
    }
    __aicore__ inline void CopyOut()
    {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        DataCopyExtParams copyParams{1, y_blockLen, 0, 0, 0};
        DataCopyPad(yGm[this->y_offset], yLocal, copyParams); 
        outQueueY.FreeTensor(yLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_Y> yGm;
    uint32_t N, H, W, C, B;
    uint32_t WAlign;
    uint16_t x_blockCount;
    uint32_t x_blockLen;
    uint32_t x_srcStride;
    uint32_t x_dstStride;
    uint32_t y_blockLen;
    uint32_t x_offset;
    uint32_t y_offset;
};
class KernelDepthToSpace_4{
public:
    __aicore__ inline KernelDepthToSpace_4() {} 
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, 
                                uint32_t blockLength,
                                uint32_t N, uint32_t H, uint32_t W, uint32_t C, uint32_t B,
                                uint32_t baseW)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->N = N;
        this->H = H;
        this->W = W;   
        this->C = C;
        this->B = B;
        xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x, blockLength);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, blockLength);
        uint32_t elemPerBlock = (ALIGN_NUM / sizeof(DTYPE_X));
        this->WAlign = RoundUp(W, elemPerBlock);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, (WAlign * (this->B) * sizeof(DTYPE_X)));
        uint32_t WBAlign = RoundUp(W*B, elemPerBlock);
        pipe.InitBuffer(outQueueY, BUFFER_NUM, (WBAlign * sizeof(DTYPE_X)));
    } 
    __aicore__ inline void Process()
    { 
        uint32_t block_id = GetBlockIdx();
        uint32_t core_num = GetBlockNum();
        this->x_blockCount = B;
        this->x_blockLen  = W * sizeof(DTYPE_X);
        this->x_dstStride = (H - 1) * W * sizeof(DTYPE_X);
        this->y_blockLen = B * W * sizeof(DTYPE_Y);
        for(int outloop = 0; outloop < (this->N); outloop++) {
            uint32_t x_out_offset = outloop * (this->C) * (this->H) * (this->W);
            for(int inloop = block_id; inloop < (((this->C) / (this->B) * (this->H))) ; inloop += core_num)
            {   
                this->x_offset = x_out_offset + ((inloop * H * B)% (C * H) + (inloop /  (B))) * W; 
                this->y_offset = x_out_offset + inloop * (B * W);
                CopyIn();
                Compute();
                CopyOut();
            }
        }
    }
private:
    __aicore__ inline uint32_t RoundUp(uint32_t a, uint32_t b)
    { 
        if (b == 0) {
            return 0;
        }
        return ((a + b - 1) / b) * b;
    }  
    __aicore__ inline void CopyIn()
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        DataCopyExtParams copyParams{x_blockCount, x_blockLen, x_dstStride, 0, 0};
        DataCopyPadExtParams<DTYPE_X> padParams{true, 0, 0, 0}; 
        DataCopyPad(xLocal, xGm[this->x_offset], copyParams, padParams); 
        inQueueX.EnQue(xLocal);
    }    
    __aicore__ inline void Compute() 
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();
        for(int i = 0; i<(this->W); i++)
        {
            for(int j=0; j<(this->B); j++)
            {
                yLocal.SetValue((i * (this->B) + j), xLocal.GetValue(j*(this->WAlign) + i));
            }
        }
        outQueueY.EnQue(yLocal);
        inQueueX.FreeTensor(xLocal);
    }
    __aicore__ inline void CopyOut()
    {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        DataCopyExtParams copyParams{1, y_blockLen, 0, 0, 0};
        DataCopyPad(yGm[this->y_offset], yLocal, copyParams); 
        outQueueY.FreeTensor(yLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_Y> yGm;
    uint32_t N, H, W, C, B;
    uint32_t WAlign;
    uint16_t x_blockCount;
    uint32_t x_blockLen;
    uint32_t x_srcStride;
    uint32_t x_dstStride;
    uint32_t y_blockLen;
    uint32_t x_offset;
    uint32_t y_offset;
};
extern "C" __global__ __aicore__ void depth_to_space(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    if(TILING_KEY_IS(1))
    {
        KernelDepthToSpace_1 OP_1;
        OP_1.Init(x, y, tiling_data.blockLength, tiling_data.N, tiling_data.H, tiling_data.W, tiling_data.C, tiling_data.B, tiling_data.baseW);
        OP_1.Process();
    }
    else if(TILING_KEY_IS(2))
    {
        KernelDepthToSpace_2 OP_2;
        OP_2.Init(x, y, tiling_data.blockLength, tiling_data.N, tiling_data.H, tiling_data.W, tiling_data.C, tiling_data.B, tiling_data.baseW);
        OP_2.Process();
    }
    else if(TILING_KEY_IS(3))
    {
        KernelDepthToSpace_3 OP_3;
        OP_3.Init(x, y, tiling_data.blockLength, tiling_data.N, tiling_data.H, tiling_data.W, tiling_data.C, tiling_data.B, tiling_data.baseW);
        OP_3.Process();
    }    
    else if(TILING_KEY_IS(4))
    {
        KernelDepthToSpace_4 OP_4;
        OP_4.Init(x, y, tiling_data.blockLength, tiling_data.N, tiling_data.H, tiling_data.W, tiling_data.C, tiling_data.B, tiling_data.baseW);
        OP_4.Process();
    }
    else if(TILING_KEY_IS(5))
    {
        KernelDepthToSpace_5 OP_5;
        OP_5.Init(x, y, tiling_data.blockLength, tiling_data.N, tiling_data.H, tiling_data.W, tiling_data.C, tiling_data.B, tiling_data.baseW);
        OP_5.Process();
    }    
}
